/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/
#include <algorithm>
#include <cstring>
#include <fstream>
#include <map>
#include <sstream>
#include <string>
#include <cctype>

#include "graph/debug/ge_op_types.h"
#include "ascir_registry.h"
namespace ge {
namespace ascir {
namespace {
const char *GetPureFileName(const char *path) {
  auto name = std::strrchr(path, '/');
  if (name == nullptr) {
    name = path;
  } else {
    ++name;
  }
  return name;
}

std::string TryGenNodeAttrClass(const AscIrDef &def, std::stringstream &ss) {
  const auto &attr_defs = def.attr_defs;
  if (attr_defs.empty()) {
    return "AscNodeAttr";
  }
  std::string derived_class_name = std::string("Asc").append(def.type).append("NodeAttr");
  // 暂时没啥好的办法，data的类需要先定义好，gen出来的话有点晚了
  if (def.type == ge::DATA) {
    ss << "  using " << derived_class_name << " = ge::" << derived_class_name << ";" << std::endl;
    return derived_class_name;
  }
  ss << "struct " << derived_class_name << ": public AscNodeAttr {" << std::endl;
  ss << "  ~" << derived_class_name << "() override = default;" << std::endl;
  ss << "  static " << derived_class_name << " &GetAttrGroup(ge::Operator &op) {" << std::endl;
  ss << "    auto opdesc = ge::OpDescUtils::GetOpDescFromOperator(op).get();\n"
        "    CHECK_NOTNULL_WITH_THROW_EXCEPTION(opdesc);" << std::endl;
  ss << "    auto attr_group = opdesc->GetOrCreateAttrsGroup<" << derived_class_name << ">();" << std::endl;
  ss << "    CHECK_NOTNULL_WITH_THROW_EXCEPTION(attr_group);\n"
        "    return *attr_group;" << std::endl;
  ss << "  };" << std::endl;

  ss << "  AttrGroupsBase *CreateNew() override {" << std::endl;
  ss << "    return new(std::nothrow) " << derived_class_name << "(*this);" << std::endl;
  ss << "  };" << std::endl;
  for (const auto &attr_def : attr_defs) {
    ss << "  " << attr_def.asc_ir_type << " " << attr_def.name << ";" << std::endl;
  }
  ss << "};" << std::endl;
  return derived_class_name;
};

void GenAscIr(const AscIrDef &def, std::stringstream &ss) {
  ss << "namespace ascir_op {" << std::endl;
  const auto &attr_defs = def.attr_defs;
  ss << "struct " << def.type << " : public ge::op::" << def.type << " {" << std::endl;
  const auto &attr_class = TryGenNodeAttrClass(def, ss);
  ss << "  " << attr_class << " &attr;" << std::endl;
  // generate input output definitions
  const auto &input_defs = def.input_defs;
  for (size_t i = 0UL; i < input_defs.size(); ++i) {
    ss << "  AscOpInput<" << i << "> " << input_defs[i].first << ";" << std::endl;
  }

  const auto &output_defs = def.output_defs;
  for (const auto &output_def : output_defs) {
    ss << "  AscOpOutput " << output_def << ";" << std::endl;
  }

  for (const auto &attr_def : attr_defs) {
    ss << "  " << attr_def.asc_ir_type << " &" << attr_def.name << ";" << std::endl;
  }

  // generate constructor func definitions
  if (def.start_node) {
    ss << "  inline " << def.type << "(const char *name, AscGraph &graph) : ge::op::"
       << def.type << "(name), attr(" << attr_class << "::GetAttrGroup(*this))";
  } else {
    ss << "  inline " << def.type << "(const char *name) : ge::op::"
       << def.type << "(name), attr(" << attr_class << "::GetAttrGroup(*this))";
  }

  for (const auto &input_def : input_defs) {
    ss << "," << std::endl << "    " << input_def.first << "(this)";
  }
  for (size_t i = 0UL; i < output_defs.size(); ++i) {
    ss << "," << std::endl << "    " << output_defs[i] << "(this, " << i << ")";
  }
  for (const auto &attr_def : attr_defs) {
    ss << "," << std::endl << "    " << attr_def.name << "(attr." << attr_def.name << ")";
  }

  if (def.start_node) {
    ss << " {" << std::endl << "    graph.AddNode(*this) ;" << std::endl << "  }" << std::endl;
  } else {
    ss << " {}" << std::endl;
  }

  ss << "};" << std::endl;
  ss << "}" << std::endl;
}

void GenGeIr(const AscIrDef &def, std::stringstream &ss) {
  ss << "REG_OP(" << def.type << ")" << std::endl;

  const auto &attr_defs = def.attr_defs;
  if (!attr_defs.empty()) {
    for (const auto &attr_def : attr_defs) {
      ss << "    .REQUIRED_ATTR(" << attr_def.name << ", " << attr_def.ge_ir_type << ")" << std::endl;
    }
  }

  const auto &input_defs = def.input_defs;
  for (const auto &input_def : input_defs) {
    if (input_def.second == ge::IrInputType::kIrInputOptional) {
      ss << "    .OPTIONAL_INPUT(" << input_def.first << ", TensorType::ALL())" << std::endl;
    } else {
      ss << "    .INPUT(" << input_def.first << ", TensorType::ALL())" << std::endl;
    }
  }

  const auto &output_defs = def.output_defs;
  for (const auto &output_def : output_defs) {
    ss << "    .OUTPUT(" << output_def << ", TensorType::ALL())" << std::endl;
  }

  ss << ".OP_END_FACTORY_REG_WITHOUT_REGISTER(" << def.type << ")" << std::endl;
}

class FunctionGenerator {
 public:
  explicit FunctionGenerator(const AscIrDef &def) : def_(def) {}
  virtual ~FunctionGenerator() = default;

  virtual void Gen(std::stringstream &ss, const bool has_optional_input) const {
    GenDefinition(ss, has_optional_input);

    GenInstantiation(ss);
    ss << std::endl;

    if (GenConnectInputs(ss, has_optional_input)) {
      ss << std::endl;
    }

    if (GenAttrAssignment(ss)) {
      ss << std::endl;
    }

    GenSchedInfo(ss);
    ss << std::endl;

    if (GenOutputsAssignment(ss)) {
      ss << std::endl;
    }
    GenOutputMemInfo(ss);
    GenPaddingAxis(ss);

    // 计算向量化轴，向量化轴的计算顺序:输出View在当前API的Loop轴内侧(不包括Loop轴)的所有轴
    TryGenOutputsVectorizedAxis(ss);

    ss << std::endl;

    GenReturn(ss);
  }

  virtual void GenDefinition(std::stringstream &ss, const bool has_optional_input) const;
  virtual void GenInstantiation(std::stringstream &ss) const;
  virtual bool GenConnectInputs(std::stringstream &ss, const bool has_optional_input) const;
  virtual bool GenAttrAssignment(std::stringstream &ss) const;
  virtual void GenSchedInfo(std::stringstream &ss) const {
    ss << "  op.attr.sched.exec_order = CodeGenUtils::GenNextExecId(op);" << std::endl;
    ss << "  SET_SCHED_AXIS_IF_IN_CONTEXT();" << std::endl;
  }
  virtual void TryGenOutputsVectorizedAxis(std::stringstream &ss) const {
    for (const auto &name : def_.output_defs) {
      ss << "  *op." << name << ".vectorized_axis = AxisUtils::GetDefaultVectorizedAxis(*op." << name
         << ".axis, op.attr.sched.loop_axis);" << std::endl;
    }
  }
  virtual void GenOutputMemInfo(std::stringstream &ss) const {
    for (const auto &name : def_.output_defs) {
      ss << "  op." << name << ".mem->tensor_id = "
         << "CodeGenUtils::GenNextTensorId(op);" << std::endl;
    }
  }
  virtual bool GenOutputsAssignment(std::stringstream &ss) const {
    bool generated = false;

    // generate infer data type code
    if (def_.infer_data_type_generator != nullptr) {
      generated = true;
      def_.infer_data_type_generator(def_, ss);
    }
    if (def_.infer_view_generator != nullptr) {
      generated = true;
      def_.infer_view_generator(def_, ss);
    }
    return generated;
  }
  virtual void GenPaddingAxis(std::stringstream &ss) const {
    for (const auto &name : def_.output_defs) {
      ss << "  THROW(PadOutputViewToSched(op." << name << "));" << std::endl;
    }
  }
  virtual void GenReturn(std::stringstream &ss) const {
    if (def_.output_defs.empty()) {
      ss << "  return op;" << std::endl;
    } else if (def_.output_defs.size() == 1U) {
      ss << "  return op." << def_.output_defs[0U] << ";" << std::endl;
    } else {
      ss << "  return std::make_tuple(";
      for (size_t i = 0; i < def_.output_defs.size(); ++i) {
        if (i == 0) {
          ss << "op." << def_.output_defs[i];
        } else {
          ss << " ,op." << def_.output_defs[i];
        }
      }
      ss << ");" << std::endl;
    }
    ss << "}" << std::endl;
  }

 protected:
  const AscIrDef &def_;
};
void ascir::FunctionGenerator::GenDefinition(std::stringstream &ss, const bool has_optional_input) const {
  const std::vector<std::pair<std::string, IrInputType>> *input_defs;
  std::vector<std::pair<std::string, IrInputType>> empty_input_defs;

  if (def_.start_node) {
    // TTODO 由于历史原因，start_node（例如Data）仍然带有输入定义，但是这种输入实际是不连边的。
    //      但是为了最小化修改，当前先不修改Data的定义，后续需要做调整，对与StartNode类型，不定义输入，
    //      或者认为没有输入的op就是start node，在定义IR时不需要再显式指定start node标记
    input_defs = &empty_input_defs;
  } else {
    input_defs = &def_.input_defs;
  }
  ss << "inline ";

  // auto OpType(const char *name, const AscOpOutput &x_in, ...) {
  ss << "auto " << def_.type << "(const char *name";
  if (!input_defs->empty()) {
    for (const auto &input_def : *input_defs) {
      if (has_optional_input || (input_def.second != ge::IrInputType::kIrInputOptional)) {
        ss << ", const ge::AscOpOutput &" << input_def.first << "_in";
      }
    }
  } else {
    ss << ", ge::AscGraph &graph";
  }

  for (const auto &attr_def : def_.attr_defs) {
    ss << ", const " << attr_def.asc_ir_type << " &" << attr_def.name;
  }
  ss << ") {" << std::endl;
}
void ascir::FunctionGenerator::GenInstantiation(std::stringstream &ss) const {
  if (def_.start_node) {
    ss << "  const auto &op_ptr = std::make_shared<ge::ascir_op::" << def_.type << ">(name, graph);" << std::endl;
  } else {
    ss << "  const auto &op_ptr = std::make_shared<ge::ascir_op::" << def_.type << ">(name);" << std::endl;
  }
  ss << "  auto &op = *op_ptr;" << std::endl;
  ss << "  const auto &desc = ge::OpDescUtils::GetOpDescFromOperator(op);" << std::endl;
  ss << "  desc->SetExtAttr(RELATED_OP, op_ptr);" << std::endl;
}
bool ascir::FunctionGenerator::GenConnectInputs(std::stringstream &ss, const bool has_optional_input) const {
  // TTODO 这里与GenFunctionDefinition同理，后续删除
  if (def_.start_node) {
    return false;
  }
  if (!def_.input_defs.empty()) {
    for (const auto &input_def : def_.input_defs) {
      if (has_optional_input || (input_def.second != ge::IrInputType::kIrInputOptional)) {
        ss << "  op." << input_def.first << " = " << input_def.first << "_in;" << std::endl;
      }
    }
  }
  return !def_.input_defs.empty();
}
bool ascir::FunctionGenerator::GenAttrAssignment(std::stringstream &ss) const {
  if (!def_.attr_defs.empty()) {
    for (const auto &attr_def : def_.attr_defs) {
      ss << "  op." << attr_def.name << " = " << attr_def.name << ';' << std::endl;
    }
  }
  return !def_.attr_defs.empty();
}

class StartNodeFuncGenerator : public FunctionGenerator {
 public:
  explicit StartNodeFuncGenerator(const AscIrDef &def) : FunctionGenerator(def) {}
  void Gen(std::stringstream &ss, const bool has_optional_input) const override {
    if (!def_.start_node || def_.output_defs.size() != 1UL) {
      return;
    }
    FunctionGenerator::Gen(ss, has_optional_input);
  }
  void GenDefinition(std::stringstream &ss, const bool has_optional_input) const override {
    // inline ascir::ops::OpType OpType
    ss << "inline "
       << "auto" << ' ' << def_.type << "(const char *name, ge::AscGraph &graph, ge::DataType dt"
       << ", const std::vector<ge::AxisId> &axes"
       << ", const std::vector<ge::expression::ExpressionPtr> &repeats"
       << ", const std::vector<ge::expression::ExpressionPtr> &strides";

    for (const auto &attr_def : def_.attr_defs) {
      ss << ", const " << attr_def.asc_ir_type << " &" << attr_def.name;
    }
    ss << ") {" << std::endl;
  }
  bool GenOutputsAssignment(std::stringstream &ss) const override {
    const auto &output_name = def_.output_defs[0];
    ss << "  op." << output_name << ".dtype = dt;" << std::endl;
    ss << "  *op." << output_name << ".axis = axes;" << std::endl;
    ss << "  *op." << output_name << ".repeats = repeats;" << std::endl;
    ss << "  *op." << output_name << ".strides = strides;" << std::endl;
    return true;
  }
};

class StoreNodeFuncGenerator : public FunctionGenerator {
 public:
  explicit StoreNodeFuncGenerator(const AscIrDef &def) : FunctionGenerator(def) {}
  void Gen(std::stringstream &ss, const bool has_optional_input) const override {
    if (def_.type != "Store") {
      return;
    }
    ss << "inline "
       << "auto" << ' ' << def_.type << "(const char *name";
    ss << ", const ge::AscOpOutput &" << "ub_in";
    ss << ", ge::AscOpOutput &gm_output";
    for (const auto &attr_def : def_.attr_defs) {
      ss << ", const " << attr_def.asc_ir_type << " &" << attr_def.name;
    }
    ss << ") {" << std::endl;
    ss << "  auto store_out = Store(name, ub_in);" << std::endl;
    ss << "  auto &gm_producer = const_cast<Operator &>(gm_output.GetOwnerOp());" << std::endl;
    ss << "  auto &store_op = const_cast<Operator &>(store_out.GetOwnerOp());" << std::endl;
    ss << "  gm_producer.SetInput(0U, store_op, 0U);" << std::endl;
    ss << "  AddEdgeForNode(store_op, 0U, gm_producer, 0U);" << std::endl;
    ss << "  auto *gm_producer_attr = CodeGenUtils::GetOwnerOpAscAttr(gm_producer);" << std::endl;
    ss << "  gm_producer_attr->sched.exec_order = CodeGenUtils::GenNextExecId(store_op);" << std::endl;
    ss << "}" << std::endl;
  }
};

class ContiguousStartNodeFuncGenerator : FunctionGenerator {
 public:
  explicit ContiguousStartNodeFuncGenerator(const AscIrDef &def) : FunctionGenerator(def) {}
  void Gen(std::stringstream &ss, const bool has_optional_input) const override {
    if (!def_.start_node || def_.output_defs.size() != 1UL) {
      return;
    }
    FunctionGenerator::Gen(ss, has_optional_input);
  }
  void GenDefinition(std::stringstream &ss, const bool has_optional_input) const override {
    ss << "inline "
       << "auto" << " Contiguous" << def_.type
       << "(const char *name, ge::AscGraph &graph, ge::DataType dt"
       << ", const std::vector<ge::Axis> &axes";

    for (const auto &attr_def : def_.attr_defs) {
      ss << ", const " << attr_def.asc_ir_type << " &" << attr_def.name;
    }
    ss << ") {" << std::endl;
  }
  bool GenOutputsAssignment(std::stringstream &ss) const override {
    const auto &output_name = def_.output_defs[0];
    ss << "  op." << output_name << ".dtype = dt;" << std::endl;
    ss << "  op." << output_name << ".SetContiguousView(axes);" << std::endl;
    return true;
  }
};

void GetHeaderGuarderFromPath(const char *path, std::stringstream &ss) {
  auto name = GetPureFileName(path);

  ss << "ASCIR_OPS_";

  while (*name != '\0') {
    auto c = toupper(*name++);
    if (c < 'A' || c > 'Z') {
      ss << '_';
    } else {
      ss << (char) c;
    }
  }

  ss << '_';
}
}  // namespace

void GenFunc(const AscIrDef &def, std::stringstream &ss) {
  FunctionGenerator(def).Gen(ss, false);
  StartNodeFuncGenerator(def).Gen(ss, false);
  ContiguousStartNodeFuncGenerator(def).Gen(ss, false);
  StoreNodeFuncGenerator(def).Gen(ss, false);
  bool has_optional_input = false;
  for (const auto &input_def : def.input_defs) {
    if (input_def.second == ge::IrInputType::kIrInputOptional) {
      has_optional_input = true;
      break;
    }
  }
  if (has_optional_input) {
    FunctionGenerator(def).Gen(ss, true);
    StartNodeFuncGenerator(def).Gen(ss, true);
    ContiguousStartNodeFuncGenerator(def).Gen(ss, true);
    StoreNodeFuncGenerator(def).Gen(ss, true);
  }
}

void GenAll(std::stringstream &ss) {
  std::stringstream ss_asc_ir;
  std::stringstream ss_ge_ir;

  ss << R"(#include "ascend_ir/ascend_ir.h")" << std::endl << std::endl;
  ss << R"(#include "ascend_ir/ascend_reg_ops.h")" << std::endl << std::endl;
  ss << R"(#include "utils/cg_utils.h")" << std::endl << std::endl;
  ss << "#include <variant>" << std::endl;
  ss << "#include <type_traits>" << std::endl;
  ss << "#include <tuple>" << std::endl << std::endl;

  std::map<std::pair<std::string, int64_t>, AscIrDef> ordered_keys_to_def;
  for (const auto &type_and_def : AscirRegistry::GetInstance().GetAll()) {
    ordered_keys_to_def[std::make_pair(type_and_def.second.file_path, type_and_def.second.line)] = type_and_def.second;
  }

  for (const auto &key_and_def : ordered_keys_to_def) {
    ss << "// Defined at " << GetPureFileName(key_and_def.second.file_path.c_str()) << ':' << key_and_def.second.line
       << std::endl;

    ss << "namespace ge {" << std::endl;
    GenGeIr(key_and_def.second, ss);
    ss << "}" << std::endl;  // namespace ge

    ss << "namespace ge {" << std::endl;
    GenAscIr(key_and_def.second, ss);
    ss << "}" << std::endl << std::endl;  // namespace ascir
  }

  ss << "namespace ge {" << std::endl;
  ss << "namespace ascir {" << std::endl;
  ss << "namespace cg {" << std::endl;
  for (auto &key_and_def : ordered_keys_to_def) {
    GenFunc(key_and_def.second, ss);
    // 如果有node属性配置，重载一个不设置属性的构造函数，把属性变成可选
    if (!key_and_def.second.attr_defs.empty()) {
      key_and_def.second.attr_defs.clear();
      GenFunc(key_and_def.second, ss);
    }
  }
  ss << "}" << std::endl;               // namespace cg
  ss << "}" << std::endl;               // namespace ascir
  ss << "}" << std::endl << std::endl;  // namespace ge
}

void GenHeaderFileToStream(const char *path, std::stringstream &ss) {
  std::stringstream ss_header_guarder;
  GetHeaderGuarderFromPath(path, ss_header_guarder);
  auto guarder = ss_header_guarder.str();

  ss << "// Generated from asc-ir definition files, "
        "any modification made to this file may be overwritten after compile."
     << std::endl;
  ss << "// If you want to add self-defined asc-ir, please create a seperated header file." << std::endl;
  ss << "#ifndef " << guarder << std::endl;
  ss << "#define " << guarder << std::endl << std::endl;

  GenAll(ss);

  ss << "#endif  // " << guarder << std::endl;
}

int GenHeaderFile(const char *path) {
  std::stringstream ss;
  GenHeaderFileToStream(path, ss);
  std::ofstream fs(path);
  if (!fs) {
    return -1;
  }
  fs << ss.str();
  fs.close();
  return 0;
}

void PreventLinkerOptimizeForAscIrGenerator() {
  GenHeaderFile("::memory::");
}
}  // namespace ascir
}

